#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
天勤账户资金和持仓信息显示工具
显示天勤账户的账户资金、可用资金、权益、默认持仓和当日盈亏等信息
"""

import tqsdk
import json
import os
from datetime import datetime
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich.layout import Layout
from rich.live import Live

# 初始化控制台输出
console = Console()

class TqAccountMonitor:
    """天勤账户监控器"""
    
    def __init__(self):
        # 默认账号信息
        self.default_account = "cps168"
        self.default_password = "alibaba"
        self.api = None
        self.account = None
        self.account_id = None
        
    def load_account_config(self):
        """从配置文件加载账号信息"""
        account = self.default_account
        password = self.default_password
        
        try:
            # 尝试从多个可能的配置文件路径加载
            config_paths = [
                os.path.join(os.path.dirname(__file__), 'db_config.json'),
                os.path.join(os.path.dirname(__file__), '../../db_config.json'),
                os.path.join(os.path.dirname(__file__), '../../app_config.json')
            ]
            
            for config_path in config_paths:
                if os.path.exists(config_path):
                    with open(config_path, 'r', encoding='utf-8') as f:
                        config = json.load(f)
                        if 'tq_username' in config and 'tq_password' in config:
                            account = config['tq_username']
                            password = config['tq_password']
                            console.print(f"从配置文件 {config_path} 加载账号信息成功", style="blue")
                            return account, password
        except Exception as e:
            console.print(f"无法从配置文件加载账号信息: {e}", style="yellow")
        
        console.print("使用默认账号信息", style="yellow")
        return account, password
    
    def login(self):
        """登录天勤账户"""
        try:
            account, password = self.load_account_config()
            
            console.print(f"正在登录天勤账号: {account}", style="blue")
            self.api = tqsdk.TqApi(account=tqsdk.TqKq(), auth=tqsdk.TqAuth(account, password))
            
            # 等待API同步数据
            console.print("等待API数据同步...", style="blue")
            for _ in range(5):
                self.api.wait_update()  # 等待数据推送
            
            # 使用账号名作为账户标识并提示
            self.account_id = account
            console.print(f"账户已连接，标识为账号名: {self.account_id}", style="green")
            return True
            
        except Exception as e:
            console.print(f"登录失败: {e}", style="red")
            if self.api:
                self.api.close()
            return False
    
    def get_account_info(self):
        """获取账户资金信息"""
        if not self.api:
            return None
        
        try:
            # 先尝试更新数据
            self.api.wait_update(0.1)
            
            # 优先使用官方 API 获取账户信息
            account_obj = None
            try:
                account_obj = self.api.get_account()
            except Exception as _:
                account_obj = None
            account_data = None
            if account_obj:
                # 将对象转为字典以兼容后续处理
                account_data = getattr(account_obj, '__dict__', None) or {
                    'available': getattr(account_obj, 'available', 0),
                    'balance': getattr(account_obj, 'balance', 0),
                    'equity': getattr(account_obj, 'equity', 0),
                    'close_profit': getattr(account_obj, 'close_profit', 0),
                    'float_profit': getattr(account_obj, 'float_profit', 0),
                    'risk_ratio': getattr(account_obj, 'risk_ratio', 0),
                    'margin': getattr(account_obj, 'margin', 0),
                    'commission': getattr(account_obj, 'commission', 0)
                }
            
            # 兼容旧路径：若仍为空再尝试 _data 路径
            if not account_data and hasattr(self.api, '_data'):
                accounts = self.api._data.get('accounts', {})
                if accounts:
                    account_data = accounts.get(list(accounts.keys())[0], {})
                elif self.account_id:
                    account_data = accounts.get(self.account_id, {})
                if not account_data:
                    trade_root = self.api._data.get('trade', {})
                    acct_key = self.account_id or (list(accounts.keys())[0] if accounts else None)
                    if isinstance(trade_root, dict) and acct_key in trade_root:
                        trade_acc = trade_root.get(acct_key, {}).get('account', {})
                        if trade_acc:
                            account_data = trade_acc
                            console.print("从 trade/account 路径获取到账户信息", style="blue")
            
            # 如果没有获取到账户数据，返回基本信息
            if not account_data:
                console.print("无法获取详细账户数据，返回基本信息", style="yellow")
                return {
                    '账户ID': self.account_id or '未知',
                    '可用资金': 0,
                    '资金权益': 0,
                    '动态权益': 0,
                    '当日盈亏': 0,
                    '浮动盈亏': 0,
                    '风险度': 0,
                    '保证金占用': 0,
                    '手续费': 0,
                    '更新时间': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                }
            
            # 提取需要的账户信息
            account_info = {
                '账户ID': self.account_id or '未知',
                '可用资金': account_data.get('available', 0),
                '资金权益': account_data.get('balance', 0),
                '动态权益': account_data.get('equity', 0),
                '当日盈亏': account_data.get('close_profit', 0),
                '浮动盈亏': account_data.get('float_profit', 0),
                '风险度': account_data.get('risk_ratio', 0),
                '保证金占用': account_data.get('margin', 0),
                '手续费': account_data.get('commission', 0),
                '更新时间': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            
            return account_info
        except Exception as e:
            console.print(f"获取账户信息失败: {e}", style="red")
            return None
    
    def get_positions(self):
        """获取账户持仓信息"""
        if not self.api:
            return []
        
        try:
            # 直接查询持仓信息（不等待同步）
            console.print("直接查询持仓信息...", style="blue")
            
            positions = []
            
            # 优先尝试使用 api.get_position() 无参获取全部持仓字典
            # 直接调用 api.get_position() 无参获取全部持仓字典（不做重试）
            try:
                pos_map = self.api.get_position()
            except Exception:
                pos_map = None
            if isinstance(pos_map, dict) and pos_map:
                for symbol, pos_obj in pos_map.items():
                    position = self._extract_position_info(symbol, pos_obj)
                    if position['多头持仓'] > 0 or position['空头持仓'] > 0:
                        positions.append(position)
                console.print(f"直接获取到 {len(positions)} 个非零持仓", style="green")
            
            # 汇总候选合约代码并使用 api.get_position 获取标准持仓对象
            candidate_symbols = set()
            try:
                accounts = self.api._data.get('accounts', {}) if hasattr(self.api, '_data') else {}
                acct_key = self.account_id or (list(accounts.keys())[0] if isinstance(accounts, dict) and accounts else None)
                if isinstance(accounts, dict) and acct_key and isinstance(accounts.get(acct_key, {}), dict):
                    pos_map = accounts.get(acct_key, {}).get('positions', {})
                    for s in (pos_map.keys() if isinstance(pos_map, dict) else []):
                        candidate_symbols.add(s)
                trade_root = self.api._data.get('trade', {}) if hasattr(self.api, '_data') else {}
                if isinstance(trade_root, dict) and acct_key in trade_root:
                    pos_map2 = trade_root.get(acct_key, {}).get('positions', {})
                    for s in (pos_map2.keys() if isinstance(pos_map2, dict) else []):
                        candidate_symbols.add(s)
                top_positions = self.api._data.get('positions', {}) if hasattr(self.api, '_data') else {}
                for s in (top_positions.keys() if isinstance(top_positions, dict) else []):
                    candidate_symbols.add(s)
            except Exception:
                pass
            
            # 对候选合约调用 get_position 并提取有效持仓
            for symbol in list(candidate_symbols):
                try:
                    pos_obj = self.api.get_position(symbol)
                    position = self._extract_position_info(symbol, pos_obj)
                    if position['多头持仓'] > 0 or position['空头持仓'] > 0:
                        positions.append(position)
                except Exception:
                    # 如果 get_position 失败，尝试从顶层原始数据解析
                    raw = (top_positions.get(symbol) if isinstance(top_positions, dict) else None)
                    position = self._extract_position_info(symbol, raw or {})
                    if position['多头持仓'] > 0 or position['空头持仓'] > 0:
                        positions.append(position)
            
            # 如果仍然没有候选合约，提示用户
            if not candidate_symbols:
                console.print("未能从 _data 路径发现任何候选合约代码", style="yellow")
            
            # 兼容回退：若仍为空，回退到 _data 路径读取持仓，并对每个 symbol 调用 api.get_position 强制解析
            if not positions and hasattr(self.api, '_data'):
                console.print("回退到 _data 路径读取持仓", style="blue")
                pos_data = {}
                accounts = self.api._data.get('accounts', {})
                acct_key = self.account_id or (list(accounts.keys())[0] if accounts else None)
                if accounts and acct_key and isinstance(accounts.get(acct_key, {}), dict):
                    pos_data = accounts.get(acct_key, {}).get('positions', {})
                if not pos_data:
                    trade_root = self.api._data.get('trade', {})
                    if isinstance(trade_root, dict) and acct_key in trade_root:
                        pos_data = trade_root.get(acct_key, {}).get('positions', {})
                if not pos_data:
                    pos_data = self.api._data.get('positions', {})
                if pos_data:
                    console.print(f"找到持仓数据，数量: {len(pos_data)}", style="green")
                else:
                    console.print("未找到持仓数据", style="yellow")
                # 对每个 symbol 调用 api.get_position 尝试获取标准 Position 对象
                for symbol, raw in pos_data.items():
                    try:
                        pos_obj = self.api.get_position(symbol)
                        position = self._extract_position_info(symbol, pos_obj)
                    except Exception:
                        position = self._extract_position_info(symbol, raw)
                    if position['多头持仓'] > 0 or position['空头持仓'] > 0:
                        positions.append(position)
            
            # 追加调试：若仍为空，输出 _data 关键路径帮助定位
            if not positions and hasattr(self.api, '_data'):
                try:
                    keys_top = list(self.api._data.keys())
                    console.print(f"_data 顶层键: {keys_top}", style="yellow")
                    trade_root = self.api._data.get('trade', {})
                    if isinstance(trade_root, dict):
                        console.print(f"trade 可用账户键: {list(trade_root.keys())}", style="yellow")
                        if self.account_id in trade_root:
                            console.print(f"trade/{self.account_id} 键: {list(trade_root.get(self.account_id, {}).keys())}", style="yellow")
                    accounts_root = self.api._data.get('accounts', {})
                    if isinstance(accounts_root, dict):
                        console.print(f"accounts 可用账户键: {list(accounts_root.keys())}", style="yellow")
                        if self.account_id in accounts_root:
                            console.print(f"accounts/{self.account_id} 键: {list(accounts_root.get(self.account_id, {}).keys())}", style="yellow")
                except Exception as e:
                    console.print(f"调试 _data 结构时出错: {e}", style="red")
            
            # 显示找到的持仓数量
            console.print(f"总共找到 {len(positions)} 个持仓", style="green")
            
            return positions
        except Exception as e:
            console.print(f"获取持仓信息失败: {e}", style="red")
            return []
    
    def _extract_position_info(self, symbol, pos_info):
        """从持仓信息中提取需要的字段"""
        # 统一使用 getattr 优先读取 Position 对象属性，兼容字典回退
        getv = lambda obj, name, default=0: (
            getattr(obj, name, None) if not isinstance(obj, dict) else obj.get(name)
        )
        # 长短仓数量
        pos_long = getv(pos_info, 'pos_long')
        pos_short = getv(pos_info, 'pos_short')
        # 兼容不同字段命名
        if not pos_long:
            pos_long = getv(pos_info, 'volume_long') or getv(pos_info, 'long_volume') or 0
        if not pos_short:
            pos_short = getv(pos_info, 'volume_short') or getv(pos_info, 'short_volume') or 0
        # 开仓均价
        open_price_long = getv(pos_info, 'open_price_long', 0) or getv(pos_info, 'price_long', 0) or getv(pos_info, 'long_price', 0)
        open_price_short = getv(pos_info, 'open_price_short', 0) or getv(pos_info, 'price_short', 0) or getv(pos_info, 'short_price', 0)
        # 浮动盈亏
        float_profit_long = getv(pos_info, 'float_profit_long', 0) or getv(pos_info, 'long_profit', 0)
        float_profit_short = getv(pos_info, 'float_profit_short', 0) or getv(pos_info, 'short_profit', 0)
        # 保证金
        margin_long = getv(pos_info, 'margin_long', 0) or getv(pos_info, 'long_margin', 0)
        margin_short = getv(pos_info, 'margin_short', 0) or getv(pos_info, 'short_margin', 0)
        # 名称
        instrument_name = getv(pos_info, 'instrument_name', 'N/A') or getv(pos_info, 'name', 'N/A')
        
        return {
            '合约代码': symbol,
            '合约名称': instrument_name if instrument_name else 'N/A',
            '多头持仓': int(pos_long) if isinstance(pos_long, (int, float)) else (int(pos_long) if str(pos_long).isdigit() else 0),
            '空头持仓': int(pos_short) if isinstance(pos_short, (int, float)) else (int(pos_short) if str(pos_short).isdigit() else 0),
            '多头均价': float(open_price_long) if isinstance(open_price_long, (int, float)) else 0.0,
            '空头均价': float(open_price_short) if isinstance(open_price_short, (int, float)) else 0.0,
            '多头浮动盈亏': float(float_profit_long) if isinstance(float_profit_long, (int, float)) else 0.0,
            '空头浮动盈亏': float(float_profit_short) if isinstance(float_profit_short, (int, float)) else 0.0,
            '多头保证金': float(margin_long) if isinstance(margin_long, (int, float)) else 0.0,
            '空头保证金': float(margin_short) if isinstance(margin_short, (int, float)) else 0.0
        }
    
    def display_account_info(self):
        """显示账户信息"""
        account_info = self.get_account_info()
        if not account_info:
            console.print(Panel("[bold red]无法获取账户信息[/]", title="账户信息"))
            return
        
        # 创建账户信息面板
        account_panel = Panel(
            Text.from_markup(
                f"[bold blue]账户ID:[/] {account_info['账户ID']}\n" +
                f"[bold green]资金权益:[/] ¥{account_info['资金权益']:,.2f}\n" +
                f"[bold green]动态权益:[/] ¥{account_info['动态权益']:,.2f}\n" +
                f"[bold green]可用资金:[/] ¥{account_info['可用资金']:,.2f}\n" +
                f"[bold {'green' if account_info['当日盈亏'] >= 0 else 'red'}]当日盈亏:[/] ¥{account_info['当日盈亏']:,.2f}\n" +
                f"[bold {'green' if account_info['浮动盈亏'] >= 0 else 'red'}]浮动盈亏:[/] ¥{account_info['浮动盈亏']:,.2f}\n" +
                f"[bold yellow]保证金占用:[/] ¥{account_info['保证金占用']:,.2f}\n" +
                f"[bold yellow]手续费:[/] ¥{account_info['手续费']:,.2f}\n" +
                f"[bold yellow]风险度:[/] {account_info['风险度']:,.2f}%\n" +
                f"[bold grey]更新时间:[/] {account_info['更新时间']}"
            ),
            title="[bold]天勤账户资金信息[/]",
            border_style="blue"
        )
        
        console.print(account_panel)
    
    def display_positions(self):
        """显示持仓信息"""
        positions = self.get_positions()
        
        if not positions:
            console.print(Panel("[bold yellow]当前无持仓[/]", title="持仓信息"))
            return
        
        # 创建持仓表格
        table = Table(title="[bold]账户持仓信息[/]", show_header=True, header_style="bold blue")
        
        table.add_column("合约代码", style="cyan")
        table.add_column("合约名称", style="magenta")
        table.add_column("多头持仓", justify="right")
        table.add_column("多头均价", justify="right")
        table.add_column("多头盈亏", justify="right", style="green")
        table.add_column("空头持仓", justify="right")
        table.add_column("空头均价", justify="right")
        table.add_column("空头盈亏", justify="right", style="red")
        
        for pos in positions:
            table.add_row(
                pos['合约代码'],
                pos['合约名称'],
                f"{pos['多头持仓']:,}",
                f"{pos['多头均价']:,.2f}",
                f"{pos['多头浮动盈亏']:,.2f}",
                f"{pos['空头持仓']:,}",
                f"{pos['空头均价']:,.2f}",
                f"{pos['空头浮动盈亏']:,.2f}"
            )
        
        console.print(table)
    
    def monitor(self):
        """启动监控"""
        try:
            # 登录账号
            if not self.login():
                return
            
            console.print("\n[bold green]===== 账户信息展示 =====[/]\n")
            
            # 显示账户信息
            self.display_account_info()
            
            # 显示持仓信息
            self.display_positions()
            
            # 提示用户如何刷新
            console.print("\n[bold yellow]提示:[/] 如需实时监控，请按Ctrl+C退出后修改代码中的监控模式。")
            
        except KeyboardInterrupt:
            console.print("\n监控已停止", style="yellow")
        except Exception as e:
            console.print(f"监控过程中发生错误: {e}", style="red")
        finally:
            if self.api:
                self.api.close()
                console.print("已关闭API连接", style="blue")

# 主函数
def main():
    console.print("[bold green]天勤账户资金和持仓信息显示工具[/]")
    console.print("[bold blue]=======================\n[/]")
    
    monitor = TqAccountMonitor()
    monitor.monitor()

if __name__ == "__main__":
    main()